from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier as KNN
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np


def main():
    iris = load_iris()
    x = iris.data[:, [0, 1]]
    y = iris.target
    # 确保 x_train 包含所有特征
    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2,
                                                        random_state=42)  # 添加 random_state 以确保结果可重复
    print("x_train shape:", x_train.shape)

    model = KNN(n_neighbors=5)
    model.fit(x_train, y_train)
    train_score = model.score(x_train, y_train)
    test_score = model.score(x_test, y_test)
    print(f'train score: {train_score:.6f};\n test score: {test_score:.6f};')

    # 可视化部分（仅选择两个特征进行可视化）
    h = .02  # step size in the mesh
    # 确保这里使用的特征索引是有效的（例如，使用前两个特征）
    x_min, x_max = x_train[:, 0].min() - 1, x_train[:, 0].max() + 1
    y_min, y_max = x_train[:, 1].min() - 1, x_train[:, 1].max() + 1
    xx, yy = np.meshgrid(np.arange(x_min, x_max, h),
                         np.arange(y_min, y_max, h))
    Z = model.predict(np.c_[xx.ravel(), yy.ravel()])
    Z = Z.reshape(xx.shape)

    plt.figure(figsize=(10, 8))
    sns.set_style("whitegrid")
    sns.scatterplot(x=x_train[:, 0], y=x_train[:, 1], hue=y_train, palette='coolwarm')
    plt.contourf(xx, yy, Z, alpha=0.2, cmap='coolwarm')
    plt.xlabel('Feature 1')
    plt.ylabel('Feature 2')
    plt.title('k = 5')
    plt.show()


if __name__ == '__main__':
    main()
